#include <config.h>
#include <Configuration.h>
#include <mpi/communicator.h>
#include <mpi/collectives.h>
#include <Utilities/OhmmsInfo.h>
#include <Utilities/RandomGenerator.h>
#include <Utilities/Timer.h>
#include <benchmark/fft_help.h>
#include <benchmark/transpose.h>
#include <fft2d/transpose2d.h>
#include <essl.h>
#include <pessl.h>
#include <fstream>

extern "C"
{
  void blacs_pinfo(int *, int *);
  void blacs_get(int *, int *, int *);
  void blacs_gridinit(int * icontxt, char* order, int* nprow, int *ncol);
  void blacs_gridinfo(int * icontxt, int* nprow, int *ncol, int *myrow, int* mycol);
//  void pdcft2(complex<double>*, complex<double>* , int*, int*, int*, double*, int*, int*);
}

inline void print_help(const string& msg)
{
  printf("%s -n fft_dim -i iterations -d \n",msg.c_str());
}

int main(int argc, char** argv)
{

  mpi::environment env(argc,argv);
  mpi::communicator mycomm;
  OhmmsInfo ohmms("fft2_omp",mycomm.rank(),0,1);

  using namespace qmcplusplus;

  int fftsize=12288;
  int niters=10;
  bool debug=false;

  int opt;
  while((opt = getopt(argc, argv, "hdn:i:")) != -1) {
    switch(opt) {
      case 'h':
        print_help("Help Message");
        return 1;
      case 'd':
        debug=true;
        break;
      case 'n':
        fftsize=atoi(optarg);
        break;
      case 'i':
        niters=atoi(optarg);
        break;
      default:
        print_help("Unknown options");
        return 1;
    }
  }

  int mypnum=-1;
  int nprocs=0;
  blacs_pinfo(&mypnum, &nprocs);

  char blacs_order='R';
  int icontext=-1;
  int nprows=1, myrow=-1;
  int npcols=mycomm.size(), mycol=-1;

  int what0=0, what1=0;
  blacs_get(&what0,&what1,&icontext);
  blacs_gridinit(&icontext,&blacs_order, &nprows, &npcols);
  blacs_gridinfo(&icontext,&nprows, &npcols,&myrow, &mycol);

  typedef complex<double> complex_type;

  //cout << " blacs " << myrow << " " << mycol << " mpi " << mycomm.rank() << " " << mycomm.size() << endl;

  int nrow_loc=fftsize/mycomm.size();
  Matrix<complex_type> in(nrow_loc,fftsize), out(nrow_loc,fftsize), in_copy(nrow_loc,fftsize);
  in=0.0;
  out=0.0;
  in_copy=0.0;
  Vector<int> ips(32);
  ips=0;

  for(int i=0; i<in.rows(); ++i)
    for(int j=0; j<in.cols(); ++j) in(i,j)=complex_type(i+mycomm.rank()*nrow_loc,j);

  pdcft2(in.data(),out.data(),fftsize,fftsize,1,1.0,icontext,ips.data());
  pdcft2(out.data(),in_copy.data(),fftsize,fftsize,-1,1.0/static_cast<double>(fftsize*fftsize),icontext,ips.data());

  //double norm=1.0;
  //if(check_array(in.data(),in_copy.data(),in_copy.size(),norm)) cout << "We are good with pdcft2 " << endl;

  Timer clock;
  clock.restart();
  for(int iter=0; iter<niters; ++iter)
  {
    pdcft2(in.data(),out.data(),fftsize,fftsize,1,1.0,icontext,ips.data());
    pdcft2(out.data(),in_copy.data(),fftsize,fftsize,-1,1.0/static_cast<double>(fftsize*fftsize),icontext,ips.data());
  }
  double dt=clock.elapsed();
  mpi::reduce(mycomm,dt);

  if(mycomm.rank()==0)
  {
    printf("PDCFT2 %d %d %d %12.4e \n", fftsize, fftsize, mycomm.size(), dt/static_cast<double>(mycomm.size()*niters));
  }

 // char fname[64];
 // sprintf(fname,"test.%d",mycol);
 // ofstream fout(fname);
 // for(int i=0; i<in.rows(); ++i)
 //   for(int j=0; j<in.cols(); ++j) fout << in(i,j)-in_copy(i,j) << endl;

  return 0;
}
